import threading
from pathlib import Path
from typing import List, Sequence

import chromadb
from django.conf import settings
from sentence_transformers import SentenceTransformer


class VectorStoreService:
    """
    管理向量数据库与嵌入模型的单例封装.
    """

    _client = None
    _model = None
    _collection_cache = {}
    _lock = threading.Lock()

    def __init__(self):
        if settings.VECTOR_DB_TYPE != 'chromadb':
            raise NotImplementedError('仅支持 chromadb 向量库, 请更新 VectorStoreService 实现')
        if VectorStoreService._client is None:
            with VectorStoreService._lock:
                if VectorStoreService._client is None:
                    Path(settings.CHROMA_PERSIST_DIRECTORY).mkdir(parents=True, exist_ok=True)
                    VectorStoreService._client = chromadb.PersistentClient(
                        path=settings.CHROMA_PERSIST_DIRECTORY
                    )

    # Embeddings
    def _get_model(self) -> SentenceTransformer:
        if VectorStoreService._model is None:
            with VectorStoreService._lock:
                if VectorStoreService._model is None:
                    VectorStoreService._model = SentenceTransformer(settings.EMBEDDING_MODEL_NAME)
        return VectorStoreService._model

    def embed(self, documents: Sequence[str]) -> List[List[float]]:
        if not documents:
            return []
        embeddings = self._get_model().encode(
            documents,
            show_progress_bar=False,
            normalize_embeddings=True
        )
        return embeddings.tolist()

    # Collections
    def get_collection_name(self, knowledge_base_id: int) -> str:
        return f'kb_{knowledge_base_id}'

    def _get_collection(self, knowledge_base_id: int):
        name = self.get_collection_name(knowledge_base_id)
        collection = VectorStoreService._collection_cache.get(name)
        if collection is None:
            with VectorStoreService._lock:
                collection = VectorStoreService._collection_cache.get(name)
                if collection is None:
                    collection = VectorStoreService._client.get_or_create_collection(name=name)
                    VectorStoreService._collection_cache[name] = collection
        return collection

    def upsert_chunks(
        self,
        knowledge_base_id: int,
        chunk_ids: Sequence[str],
        documents: Sequence[str],
        metadatas: Sequence[dict]
    ):
        if not chunk_ids:
            return
        collection = self._get_collection(knowledge_base_id)
        embeddings = self.embed(documents)
        collection.upsert(
            ids=list(chunk_ids),
            documents=list(documents),
            metadatas=list(metadatas),
            embeddings=embeddings
        )

    def delete_vectors(self, knowledge_base_id: int, vector_ids: Sequence[str]):
        if not vector_ids:
            return
        collection = self._get_collection(knowledge_base_id)
        collection.delete(ids=list(vector_ids))

    def query(self, knowledge_base_id: int, text: str, top_k: int = 3):
        if top_k <= 0:
            return []
        collection = self._get_collection(knowledge_base_id)
        response = collection.query(
            query_texts=[text],
            n_results=top_k,
            include=['metadatas', 'documents', 'distances']
        )

        ids = response.get('ids', [[]])
        documents = response.get('documents', [[]])
        metadatas = response.get('metadatas', [[]])
        distances = response.get('distances', [[]])

        if not ids or not ids[0]:
            return []

        primary_ids = ids[0]
        doc_list = documents[0] if documents else []
        metadata_list = metadatas[0] if metadatas else []
        distance_list = distances[0] if distances else []

        results = []
        for idx, vector_id in enumerate(primary_ids):
            doc_text = doc_list[idx] if idx < len(doc_list) else ''
            metadata = metadata_list[idx] if idx < len(metadata_list) else {}
            distance = distance_list[idx] if idx < len(distance_list) else None
            results.append({
                'vector_id': vector_id,
                'content': doc_text,
                'metadata': metadata or {},
                'distance': distance,
            })
        return results
